import base64
import json

from Cryptodome.Cipher import PKCS1_OAEP
from Cryptodome.Signature import pkcs1_15
from gmssl import sm3


# 生成签名
def generate_signature(private_key, data_str):
    # 计算SM3哈希值
    sm3_hash = sm3.sm3_hash(data_str.encode('utf-8'))
    # 使用私钥进行签名
    signer = pkcs1_15.new(private_key)
    signature = signer.sign(sm3_hash.hex())

    # 返回base64编码后的签名结果
    return base64.b64encode(signature).decode('utf-8')


# 生成加密数据
def generate_encrypted_data(public_key, data):
    # 将数据转换为字节流
    data_bytes = data.encode('utf-8')

    # 使用公钥进行加密
    cipher = PKCS1_OAEP.new(public_key)
    encrypted_data = cipher.encrypt(data_bytes)

    # 返回base64编码后的加密结果
    return base64.b64encode(encrypted_data).decode('utf-8')


# 构建报文头
def build_header(appid, timestamp, sign, verify, private_key,secret):
    # 构建待签名数据
    data_str = f"appid={appid}&secret={secret}&sign={sign}&timestamp={timestamp}"

    # 生成签名
    signature = generate_signature(private_key, data_str)

    # 构建报文头
    header = {
        "appid": appid,
        "timestamp": timestamp,
        "sign": sign,
        "apisign": signature,
        "verify": verify,
    }

    return header


# 加密报文体
def encrypt_payload(public_key, payload):
    # 将报文体转换为JSON字符串
    payload_str = json.dumps(payload)

    # 生成加密数据
    encrypted_data = generate_encrypted_data(public_key, payload_str)

    # 构建加密后的报文体
    encrypted_payload = {
        "data": encrypted_data
    }

    return encrypted_payload


# 测试类
def test():
    # 加载私钥和公钥
    private_key_str = "hrgpTGgIUTb726OsIfZUOju/2akNZFHQOjtWYGfrtQY="
    public_key_str = "BK95GNW/n6K467X9atPB+cbJzi38G7yXo/kug6fLK3KeY6DFrtZuE5sj/imDZm0p47UCG+ydRvmqGMQNQdke5O8="
    private_key = (base64.b64decode(private_key_str))
    public_key = (base64.b64decode(public_key_str))

    # 构建报文体
    payload = {
        "regulatorAccNo": "755937428810701",
        "merchantId": "MerchID20210805103419610vIjJaa"
    }

    # 构建报文头
    header = build_header(
        appid="hksup004",
        timestamp="1638947960958",
        sign="telecom",
        verify="SM2",
        private_key=private_key,
        secret="secret"
    )

    # 加密报文体
    encrypted_payload = encrypt_payload(public_key, payload)

    # 打印报文头和加密后的报文体
    print("Header:", header)
    print("Encrypted Payload:", encrypted_payload)


if __name__ == '__main__':
    test()
